import re
import sys
from collections import defaultdict
from contextlib import nullcontext
from io import StringIO
from multiprocessing import Process, Queue
from typing import List, Optional, Type, Union

from filelock import FileLock
from timeout_decorator import timeout as tm

from ..schema import ActionReturn, ActionStatusCode
from .base_action import BaseAction
from .parser import BaseParser, JsonParser


class IPythonProcess(Process):
    """
    IPython进程类
    
    用于在独立进程中运行IPython代码，支持超时控制和会话管理
    """

    def __init__(self,
                 in_q: Queue,
                 out_q: Queue,
                 timeout: int = 20,
                 ci_lock: str = None,
                 daemon: bool = True):
        """
        初始化IPython进程
        
        参数:
            in_q: 输入队列，用于接收代码
            out_q: 输出队列，用于返回结果
            timeout: 代码执行超时时间（秒）
            ci_lock: 文件锁路径，用于避免核心损坏
            daemon: 是否为守护进程
        """
        super().__init__(daemon=daemon)
        self.in_q = in_q
        self.out_q = out_q
        self.timeout = timeout
        # 会话ID到shell实例的映射
        self.session_id2shell = defaultdict(self.create_shell)
        # 创建文件锁或空上下文管理器
        self.ci_lock = FileLock(ci_lock) if ci_lock else nullcontext()
        # 用于移除ANSI转义序列的正则表达式
        self._highlighting = re.compile(r'\x1b\[\d{,3}(;\d{,3}){,3}m')

    def run(self):
        """
        进程主循环
        
        处理输入队列中的消息，执行代码并返回结果
        """
        while True:
            msg = self.in_q.get()
            if msg == 'reset':  # 重置所有shell
                for session_id, shell in self.session_id2shell.items():
                    with self.ci_lock:
                        try:
                            shell.reset(new_session=False)
                        except Exception:
                            # 如果重置失败，创建新的shell实例
                            self.session_id2shell[session_id] = self.create_shell()
                self.out_q.put('ok')
            elif isinstance(msg, tuple) and len(msg) == 3:
                # 处理代码执行请求
                i, session_id, code = msg
                res = self.exec(session_id, code)
                self.out_q.put((i, session_id, res))

    def exec(self, session_id, code):
        """
        执行代码
        
        在指定会话的shell中执行代码，支持超时控制
        
        参数:
            session_id: 会话ID
            code: 要执行的代码
            
        返回:
            dict: 执行结果，包含状态和输出
        """
        try:
            shell = self.session_id2shell[session_id]
            with StringIO() as io:
                # 重定向标准输出
                old_stdout = sys.stdout
                sys.stdout = io
                
                # 执行代码（带或不带超时控制）
                if self.timeout is False or self.timeout < 0:
                    shell.run_cell(self.extract_code(code))
                else:
                    tm(self.timeout)(shell.run_cell)(self.extract_code(code))
                    
                # 恢复标准输出并获取结果
                sys.stdout = old_stdout
                output = self._highlighting.sub('', io.getvalue().strip())
                output = re.sub(r'^Out\[\d+\]: ', '', output)
                
            # 处理错误情况
            if 'Error' in output or 'Traceback' in output:
                output = output.lstrip('-').strip()
                if output.startswith('TimeoutError'):
                    output = 'The code interpreter encountered a timeout error.'
                return {'status': 'FAILURE', 'msg': output, 'code': code}
            return {'status': 'SUCCESS', 'value': output, 'code': code}
        except Exception as e:
            return {'status': 'FAILURE', 'msg': str(e), 'code': code}

    @staticmethod
    def create_shell(enable_history: bool = False, in_memory: bool = True):
        """
        创建新的IPython shell实例
        
        参数:
            enable_history: 是否启用历史记录
            in_memory: 是否使用内存存储历史记录
        """
        from IPython import InteractiveShell
        from traitlets.config import Config

        c = Config()
        c.HistoryManager.enabled = enable_history
        if in_memory:
            c.HistoryManager.hist_file = ':memory:'
        shell = InteractiveShell(config=c)
        return shell

    @staticmethod
    def extract_code(text: str) -> str:
        """
        从标记语言中提取Python代码
        
        支持从Markdown代码块或JSON中提取代码
        
        参数:
            text: 包含代码的文本
            
        返回:
            str: 提取的Python代码
        """
        import json5

        # 首先匹配三重反引号代码块
        triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
        # 然后匹配单反引号代码块
        single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
        
        if triple_match:
            text = triple_match.group(1)
        elif single_match:
            text = single_match.group(1)
        else:
            try:
                text = json5.loads(text)['code']
            except Exception:
                pass
        return text


class IPythonInteractiveManager(BaseAction):
    """
    IPython交互管理器
    
    管理多个IPython进程，支持并行代码执行
    """

    def __init__(
        self,
        max_workers: int = 50,
        timeout: int = 20,
        ci_lock: str = None,
        description: Optional[dict] = None,
        parser: Type[BaseParser] = JsonParser,
    ):
        """
        初始化管理器
        
        参数:
            max_workers: 最大工作进程数
            timeout: 代码执行超时时间
            ci_lock: 文件锁路径
            description: 动作描述
            parser: 参数解析器类
        """
        super().__init__(description, parser)
        self.max_workers = max_workers
        self.timeout = timeout
        self.ci_lock = ci_lock
        self.id2queue = defaultdict(Queue)  # 进程ID到输入队列的映射
        self.id2process = {}  # 进程ID到进程实例的映射
        self.out_queue = Queue()  # 统一的输出队列

    def __call__(self, commands: Union[str, List[str]], session_ids: Union[int, List[int]] = None):
        """
        执行代码
        
        支持批量执行多个代码块
        """
        # 标准化输入参数
        if isinstance(commands, list):
            batch_size = len(commands)
            is_batch = True
        else:
            batch_size = 1
            commands = [commands]
            is_batch = False
            
        # 处理会话ID
        if session_ids is None:
            session_ids = range(batch_size)
        elif isinstance(session_ids, int):
            session_ids = [session_ids]
            
        # 验证参数
        if len(session_ids) != batch_size or len(session_ids) != len(set(session_ids)):
            raise ValueError('the size of `session_ids` must equal that of `commands`')
            
        try:
            # 执行代码块
            exec_results = self.run_code_blocks([
                (session_id, command)
                for session_id, command in zip(session_ids, commands)
            ])
        except KeyboardInterrupt:
            self.clear()
            exit(1)
            
        # 处理执行结果
        action_returns = []
        for result, code in zip(exec_results, commands):
            action_return = ActionReturn({'command': code}, type=self.name)
            if result['status'] == 'SUCCESS':
                action_return.result = [
                    dict(type='text', content=result['value'])
                ]
                action_return.state = ActionStatusCode.SUCCESS
            else:
                action_return.errmsg = result['msg']
                action_return.state = ActionStatusCode.API_ERROR
            action_returns.append(action_return)
            
        # 返回结果
        if not is_batch:
            return action_returns[0]
        return action_returns

    def process_code(self, index, session_id, code):
        """处理单个代码块的执行"""
        ipy_id = session_id % self.max_workers
        input_queue = self.id2queue[ipy_id]
        
        # 获取或创建进程
        proc = self.id2process.setdefault(
            ipy_id,
            IPythonProcess(
                input_queue,
                self.out_queue,
                self.timeout,
                self.ci_lock,
                daemon=True))
                
        # 确保进程在运行
        if not proc.is_alive():
            proc.start()
            
        # 发送代码到进程
        input_queue.put((index, session_id, code))

    def run_code_blocks(self, session_code_pairs):
        """
        运行多个代码块
        
        并行执行多个代码块，并收集结果
        """
        size = len(session_code_pairs)
        # 分发代码到进程
        for index, (session_id, code) in enumerate(session_code_pairs):
            self.process_code(index, session_id, code)
            
        # 收集结果
        results = []
        while len(results) < size:
            msg = self.out_queue.get()
            if isinstance(msg, tuple) and len(msg) == 3:
                index, _, result = msg
                results.append((index, result))
        results.sort()
        return [item[1] for item in results]

    def clear(self):
        """清理所有资源"""
        self.id2queue.clear()
        for proc in self.id2process.values():
            proc.terminate()
        self.id2process.clear()
        while not self.out_queue.empty():
            self.out_queue.get()

    def reset(self):
        """重置所有shell实例"""
        cnt = 0
        for q in self.id2queue.values():
            q.put('reset')
            cnt += 1
        while cnt > 0:
            msg = self.out_queue.get()
            if msg == 'ok':
                cnt -= 1
